import paddle.fluid.layers as L
from paddle.fluid.dygraph import Dropout
from helm.static.models.layers import Conv2d, Pool2d, Linear, Layer, Sequential


__all__ = ["LeNet5"]


class LeNet5(Layer):

    def __init__(self, in_channels=1, num_classes=10, dropout=None):
        super().__init__()
        self.features = Sequential(
            Conv2d(in_channels, 6, 5, stride=1, padding=0, act='default'),
            Pool2d(2, 2, type='max'),
            Conv2d(6, 16, 5, stride=1, padding=0, act='default'),
            Pool2d(2, 2, type='max'),
        )

        fc = [
            Linear(400, 120, act='default'),
            Linear(120, 84, act='default'),
            Linear(84, num_classes)
        ]
        if dropout:
            fc.insert(0, Dropout(dropout, dropout_implementation='upscale_in_train'))
            fc.insert(2, Dropout(dropout, dropout_implementation='upscale_in_train'))
            fc.insert(4, Dropout(dropout, dropout_implementation='upscale_in_train'))
        self.fc = Sequential(*fc)

    def forward(self, x):
        x = self.features(x)
        x = L.flatten(x, 1)
        x = self.fc(x)
        return x